import mindspore as ms
from mindspore import context as con
from mindspore import nn, dataset, ops
from sklearn.model_selection import train_test_split
from mindvision.engine.callback import LossMonitor

import numpy as np

con.set_context(device_target='CPU', mode=con.PYNATIVE_MODE)


def sample_data(ratio=100, test_size=0.3):
    x = np.linspace(-10, 10, 21 * ratio).astype(np.float32)
    y = 3*x + 4 + np.random.rand(21 * ratio).astype(np.float32) /100
    X_train, X_test, y_train, y_test = train_test_split(
        x, y, test_size=test_size,shuffle=True)

    return [X_train, y_train], [X_test, y_test]


class MySets:
    def __init__(self, *inputs) -> None:
        self.x = inputs[0].reshape(-1,1)
        self.y = inputs[1].reshape(-1,1)

    def __len__(self):
        return len(self.x)

    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]


train_set, test_set = sample_data()

sets = {"train": dataset.GeneratorDataset(MySets(*train_set), ['data', 'target']).batch(256),
        "test": dataset.GeneratorDataset(MySets(*test_set), ['data', 'target']).batch(512)
        }


class Net(nn.Cell):
    def __init__(self, in_dim=1, auto_prefix=True, flags=None):
        super().__init__(auto_prefix, flags)

        self.fc = nn.Dense(in_dim, 1)

    def construct(self, *inputs, **kwargs):
        return self.fc(*inputs)


net = Net()
Loss_item = nn.MSELoss()
optim = nn.optim.SGD(net.trainable_params(), learning_rate=1e-3, momentum=0.9)


model = ms.Model(net, Loss_item, optim, metrics={
                 "MSE": nn.MSE(), "MAE": nn.MAE()})

model.train(100, sets['train'], callbacks=LossMonitor(1e-3, 10))

out = model.eval(sets['test'])
print(out)


# loss_net = nn.WithLossCell(net, Loss_item)

# TrainNet = nn.TrainOneStepCell(loss_net,optim)

# if __name__ == "__main__":

#     for _ in range(100):
#         loss_value = 0.0
#         count = 0
#         for x,y in sets:
#             loss_value += TrainNet(x,y)

#             if count % 10 == 9:
#                 print("Loss is {:4.6f}".format(loss_value.asnumpy()/10))
#                 loss_value = 0.0

#             count += 1
